#include "spm.h"
#include <base/utility.h>
#include <base/common.h>
#include <base/km_matcher.hpp>
#include <base/parallel_utility.hpp>
#include <output/output.h>
#include "utils.h"

using namespace SPM;

SignPoleMatching::SignPoleMatching() : debug_message_(false)
{

}

bool SignPoleMatching::Calculate_Pose(const DlFrame &dl,
    const std::vector<hdmap_sign>& hdmap_signs,
    Pose &pose,
    Camera &cam, 
    double &rmse, 
    DbgFrame *dbg_frame) 
{
    if (dl.signs.empty() || hdmap_signs.empty())
    {
        return false;
    }

    geometric_CRITERIA_ = (float)cam.width / 6.0;
    std::vector<std::pair<hdmap_sign, envelope2D>> proj_signs = ProjectHDMap(cam, pose, hdmap_signs);
    const size_t hdsign_num = proj_signs.size();
    const size_t dlsign_num = dl.signs.size();

    std::vector<int> hd_match_pair(hdsign_num, -1);
    std::vector<int> dl_match_pair(dlsign_num, -1);

    /* calcuate the matching cost */
    std::vector<std::vector<float>> KM_cost(hdsign_num, std::vector<float>(dlsign_num, 0));
    auto get_KM_data = [&KM_cost](unsigned hd_idx, unsigned dl_idx) 
    {
		return KM_cost[hd_idx][dl_idx];
	};

    cal_KM_cost_parallel(KM_cost, proj_signs, dl.signs);
    auto get_KM = std::bind(get_KM_data, std::placeholders::_1, std::placeholders::_2);
	algo_utils::MunkresData<float>	KM_matching(hdsign_num, dlsign_num, get_KM);
	auto matches = KM_matching.solve();
    
    cv::Scalar colors[6] = {cv::Scalar(255, 0, 0), cv::Scalar(0, 255, 0), cv::Scalar(0, 0, 255), cv::Scalar(0, 255, 255), cv::Scalar(34, 56, 155), cv::Scalar(130, 20, 20)};
    if (debug_message_) { std::cout << "------ Matches Info ------" << std::endl; }
    for (std::size_t i = 0; i < matches.size(); i++)
    {
        const auto &match = matches[i];
        auto &proj_sign = proj_signs[match.first];
        auto &dl_sign = dl.signs[match.second];

        if (KM_cost[match.first][match.second] != geometric_CRITERIA_)
        {
            hd_match_pair[match.first] = match.second;
            dl_match_pair[match.second] = match.first;
            if (debug_message_)
            {
                //dbg_frame->DrawRectangle(proj_signs[match.first].second, colors[i % 6], 2.0);
               // dbg_frame->DrawRectangle(dl.signs[match.second].calibrated_bbox, colors[i % 6], 2.0);
                std::cout << "  hdSign_ " << proj_sign.first.id << " matches dlSign_ " << dl_sign.id << std::endl;
            }
        }
        else
        {
            if (debug_message_)
            {
                std::cout << "  hdSign_" << proj_sign.first.id << " failed geometric check of dlSign_ " << dl_sign.id  << std::endl;
            }
        }
    }
    if (debug_message_)
    {
        std::cout << "--------------------------" << std::endl
                  << std::endl;
    }

    std::vector<std::pair<DlSign, hdmap_sign>> match_results;
    for(int hd_idx = 0; hd_idx < hd_match_pair.size(); hd_idx++)
    {
        if(hd_match_pair[hd_idx] != -1)
        {
            DlSign dlSign = dl.signs[hd_match_pair[hd_idx]];
            hdmap_sign hdSign = proj_signs[hd_idx].first;
            match_results.emplace_back(dlSign, hdSign);
            if (debug_message_)
            {
                dbg_frame->DrawRectangle2D(dlSign.calibrated_bbox, colors[hd_idx % 6], 2.0);
                dbg_frame->DrawRectangle2D(proj_signs[hd_idx].second, colors[hd_idx % 6], 2.0);
            }
           
        }
    }

    Eigen::Matrix3d R = Eigen::Matrix3d::Identity();
    Eigen::Vector3d T = Eigen::Vector3d::Zero();

    std::vector<std::pair<Eigen::Vector2d, Eigen::Vector3d>> pnp_results;
    if (utils::ComputeRT(match_results, pose, cam, R, T, rmse, 0.5, pnp_results))
    {
        Eigen::Matrix3d spmR = pose.orientation.toRotationMatrix() * R.transpose();
        Eigen::Vector3d spmT = pose.position - pose.orientation.toRotationMatrix() * R.transpose() * T;

        pose.position.x() = spmT.x();
        pose.position.y() = spmT.y();
        pose.position.z() = spmT.z();
        pose.orientation = Eigen::Quaternion<double>(spmR);
        return true;
    }

    return false;
}


std::vector<std::pair<hdmap_sign, envelope2D>> SignPoleMatching::ProjectHDMap(const Camera &cam, const Pose& pose, const std::vector<hdmap_sign>& hdmap_signs)
{
	Eigen::Matrix3d bRw = pose.orientation.toRotationMatrix();
    Eigen::Vector3d bTw = pose.position;
    Eigen::Matrix3d cRb = cam.mount.rotation.matrix();
    Eigen::Vector3d cTb = cam.mount.translation;

	std::vector<std::pair<hdmap_sign, envelope2D>> sign_corners_pair; 
	for (hdmap_sign sign : hdmap_signs)
	{
		std::vector<Eigen::Vector2d> sign_corners;
		for (std::size_t i = 0; i < sign.geometry.size(); i++)
		{
			Eigen::Vector2d uv;
			Eigen::Vector3d pcd = cRb.transpose() * (bRw.transpose() * (sign.geometry[i] - bTw) - cTb);
			if (cam.Project3d_2d(pcd, uv))
			{
				sign_corners.emplace_back(uv.x(), uv.y());
			}
		}
		envelope2D sign_env = envelope2D::GetEnvelope(sign_corners);
		sign_corners_pair.emplace_back(sign, sign_env);
	}
	return sign_corners_pair;
}

void SignPoleMatching::cal_KM_cost_parallel(std::vector<std::vector<float>> &_KM_cost, const std::vector<std::pair<hdmap_sign, envelope2D>>& hd_signs, const std::vector<DlSign>& dl_signs)
{
    size_t sign_num = _KM_cost.size();
    // calculate the cos of each pair
    for (size_t hd_idx = 0; hd_idx < sign_num; ++hd_idx)
    {
        auto &hd_sign = hd_signs[hd_idx];
        auto do_work = [this, &hd_sign, &hd_idx, &_KM_cost, &dl_signs](size_t dl_idx, std::mutex &_m) {
            auto &dl_sign = dl_signs[dl_idx];
            float iou = envelope2D::IOU(hd_sign.second, dl_signs[dl_idx].calibrated_bbox);
            float geometrical_cost = envelope2D::CenterDiff(hd_sign.second, dl_sign.calibrated_bbox);
            if (geometrical_cost > geometric_CRITERIA_ || iou < 0.3)
            {
                _KM_cost[hd_idx][dl_idx] = geometric_CRITERIA_;
            }
            else
            {
                _KM_cost[hd_idx][dl_idx] = geometrical_cost;
            }
        };

        /* do_work parallelly */
        size_t chunk = LOOP_PARALISM::get_chunk(dl_signs.size());
        for (size_t i = 0; i < chunk; ++i)
        {
            LOOP_PARALISM loop_paraller;
            auto &j = loop_paraller.thread_index_;
            for (j = 0; j < loop_paraller.numthreads_; ++j)
            {
                auto &&index = i * chunk + j;
                if (index >= dl_signs.size())
                {
                    continue;
                }
                loop_paraller.works_[j] = std::thread(do_work, index, std::ref(loop_paraller.m_));
            }
        }
    }
}